import json
import re, os
import traceback
import time
from platform import system

from event_handle.pipeline import EventHandlePipeline
from util import llm
from util.common import remove_markdown_code_block

MODEL_NAME = os.getenv('MODEL_NAME', "deepseek-chat")
# SYSTEM_PROMPT = """人设信息：
# 你是一个优秀的信息安全专家，可以根据安全问题描述和违规依据给出合理的建议
# 输入：
# 输入分为两部分：安全问题描述、违规依据
# 安全问题描述为安全检测系统给出的安全事件名称、发生时间、事件信息；违规依据为安全相关法律及规定条文
# 工作流：
# 1. 根据输入违规描述，在违规依据中查找相关条文，并列出
# 2. 根据输入违规描述及相关规定条文，给出处置建议。处置建议要求分条列出
# 用户数据：
# 安全问题描述：
# {}
# 违规依据：
# {}
# """

SYSTEM_PROMPT = """人设信息：
你是一个优秀的信息安全专家，可以在用户的安全描述中提取出精准的实体信息。实体信息包括：
ip地址：描述中可能包含了ip地址
用户名：描述中可能包含了用户名
url：描述中可能包含了url地址
密码信息：描述中可能包含了明文密码、弱密码、使用url传递密码、接口返回内容中包含密码
配置代码文件：描述中可能包含了配置文件、代码相关文件、git文件、svn文件

工作流：
1. 根据输入描述，在提取相关实体，实体类别包含上面6种。
2. 列出实体类别，使用json object格式，具体格式由【输出格式】部分给出。请严格按照格式输出

输出格式：
输出格式为json object，每个元素表示一个实体，元素的key为实体类别，value为实体值。
输出json的key必须在[ip地址、用户名、url、密码信息、配置代码文件]这个集合中
请严格按json格式输出，**不要**输出markdown格式。


样例：
如：用户描述为"检测到零信任高频访问，统计结果共命中449条日志： [访问次数：5647次，用户名：111@asi**o.com，请求url：10.22.22.123:8080]，触发条件'访问次数' > 1500"出现了ip地址、用户名，则返回'{{"ip地址":"10.22.22.123:8080","用户名":"111@asi**o.com"}}'
用户数据：
安全问题描述：
{}
"""

EVENT_HANDLE_MAP = {
    "ip地址": "发现描述中有ip地址【{}】，建议排查此地址是否为公司内部ip，此ip是否有其他可疑行为",
    "用户名": "发现描述中有用户名信息【{}】，建议定位此用户的操作日志，检查是否为合法用户（如公司内网用户、合作伙伴用户登），确定是否违规",
    "url": "发现描述中有url【{}】，建议查看url是否为公司内部url，本用户是否有访问权限",
    "密码信息": "发现描述中有密码信息【{}】，建议查看接口中是否包含用户密码信息，是否为安全测试行为",
    "配置代码文件": "发现描述中有配置、代码文件，建议排查源码仓库里是否存在源码泄露，包括git、svn、配置文件等"
}

POLISH_SYSTEM_PROMPT = """你是一个事件处置专家，会根据已有信息给出处置建议。
以下我将给出若干可疑实体，请根据我给定的格式将其润色，输出一篇处置建议。处置建议的第二部分为“整改措施”，请根据给定的可疑实体给出整改建议，如：遇到ip类问题，则给出防火墙相关的措施；遇到账号相关问题，建议用户限制账号登录等。
输出格式：
一、排除误报
{}
二、整改措施：

可疑实体：
{}
"""


async def get_entity(alarm_desc, violation_detail, tags):
    # print(violation_detail)
    # system_prompt = SYSTEM_PROMPT.format(alarm_desc, violation_detail)
    system_prompt = SYSTEM_PROMPT.format(alarm_desc + "\n" + tags)
    result = await llm.http_request(
        system_prompt,
        history=[],
        prompt='请根据已知信息，提取相应实体。务必遵循输出格式的要求。',
        # model="deepseek-v3",
        # model="deepseek-r1-distill-qwen-1.5b"
        # model="hunyuan-lite",
        model=MODEL_NAME
    )
    result = remove_markdown_code_block(result)
    try:
        result = json.loads(result)
        tmp = {}
        for k, v in result.items():
            if v and v not in ["无", "未知"]:
                tmp[k] = v
        result = tmp
    except Exception as e:
        print(traceback.format_exc())
        result = []
    print(result)
    return result


async def get_llm_result(alarm_desc, violation_detail, tags, stream=False):
    a = time.time()
    entity = await get_entity(alarm_desc, violation_detail, tags)
    b = time.time()
    print("get_entity cost: {}".format(b - a))

    event_list = []
    entity_list = []
    for e in entity:
        if e in EVENT_HANDLE_MAP:
            event_list.append(EVENT_HANDLE_MAP.get(e, "").format(entity[e]))
            entity_list.append("{}: {}".format(e, entity[e]))

    # return  list(event_list)
    async def func(s_list: list[str]):
        for x in s_list:
            yield x

    async def gen_combiner(a, b):
        async for x in a:
            yield x
        async for x in b:
            yield x

    if len(entity) == 0:
        msg = "但在信息中未能发现相关实体"
    else:
        msg = "这里有几个关键信息点引起了我的注意：\n{}".format("\n".join(entity_list))
    think = ["<think>", "用户给了问题描述信息，我需要从中提取出可能有问题的实体，并给出建议。\n{}".format(msg), "</think>"]
    if not event_list:
        solution = list("暂未发现问题，请根据日志排查")
        return func(think + solution)

    entity_prompt = "\n".join(["{}. {}".format(i + 1, x) for i, x in enumerate(event_list)])
    return gen_combiner(
        func(think),
        llm.sse_request(
            system_prompt=POLISH_SYSTEM_PROMPT.format(entity_prompt, ";".join(event_list)),
            history=[],
            prompt="请根据问题回答",
            model=MODEL_NAME
            # model="hunyuan-turbo"
        )
    )

    # return func(think + list("；".join(event_list)))


if __name__ == "__main__":
    # s = get_llm_result(
    #     "经安全监测系统检测，中国联通吉林省分公司发生【密码透出】安全事件。该事件首次发现于 2025 年 04 月 23 日 14 时 10 分 23 秒，在接口返回内容中发现密码信息，攻击者可在传输过程中进行监听或拦截，从而获取密码信息。",
    #     "")
    # print(s)
    # print("===========")
    s = SYSTEM_PROMPT.format(
        "aa")
    print(s)
